// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/vector/vector_saver.h>
#include <fstream>
#include <pollux/vector/complex_vector.h>
#include <pollux/vector/flat_vector.h>

namespace kumo::pollux {
    namespace {
        enum class Encoding : int8_t {
            kFlat = 0,
            kConstant = 1,
            kDictionary = 2,
            kLazy = 3,
        };

        template<typename T>
        void write(const T &value, std::ostream &out) {
            out.write(reinterpret_cast<const char *>(&value), sizeof(T));
        }

        template<>
        void write<bool>(const bool &value, std::ostream &out) {
            write<int8_t>(value ? 1 : 0, out);
        }

        template<>
        void write<std::string>(const std::string &value, std::ostream &out) {
            write<int32_t>(value.size(), out);
            out.write(value.data(), value.size());
        }

        template<typename T>
        T read(std::istream &in) {
            T value;
            in.read(reinterpret_cast<char *>(&value), sizeof(T));
            return value;
        }

        template<>
        bool read<bool>(std::istream &in) {
            int8_t value = read<int8_t>(in);
            return value != 0;
        }

        template<>
        std::string read<std::string>(std::istream &in) {
            auto size = read<int32_t>(in);
            std::string data;
            data.resize(size);
            in.read(data.data(), size);
            return data;
        }

        void writeEncoding(VectorEncoding::Simple encoding, std::ostream &out) {
            switch (encoding) {
                case VectorEncoding::Simple::FLAT:
                case VectorEncoding::Simple::ROW:
                case VectorEncoding::Simple::ARRAY:
                case VectorEncoding::Simple::MAP:
                    write<int32_t>(static_cast<int8_t>(Encoding::kFlat), out);
                    return;
                case VectorEncoding::Simple::CONSTANT:
                    write<int32_t>(static_cast<int8_t>(Encoding::kConstant), out);
                    return;
                case VectorEncoding::Simple::DICTIONARY:
                    write<int32_t>(static_cast<int8_t>(Encoding::kDictionary), out);
                    return;
                case VectorEncoding::Simple::LAZY:
                    write<int32_t>(static_cast<int8_t>(Encoding::kLazy), out);
                    return;
                default:
                    POLLUX_UNSUPPORTED("Unsupported encoding: {}", mapSimpleToName(encoding));
            }
        }

        Encoding readEncoding(std::istream &in) {
            auto encoding = Encoding(read<int32_t>(in));
            switch (encoding) {
                case Encoding::kFlat:
                case Encoding::kConstant:
                case Encoding::kDictionary:
                case Encoding::kLazy:
                    return encoding;
                default:
                    POLLUX_UNSUPPORTED("Unsupported encoding: {}", encoding);
            }
        }

        /// Serializes a BufferPtr into binary format and writes it to the
        /// provided output stream. 'buffer' must be non-null.
        void writeBuffer(const BufferPtr &buffer, std::ostream &out) {
            write<int32_t>(buffer->size(), out);
            out.write(buffer->as<char>(), buffer->size());
        }

        /// Serializes a optional BufferPtr into binary format and writes it to the
        /// provided output stream.
        void writeOptionalBuffer(const BufferPtr &buffer, std::ostream &out) {
            if (buffer) {
                write<bool>(true, out);
                writeBuffer(buffer, out);
            } else {
                write<bool>(false, out);
            }
        }

        /// Deserializes a BufferPtr serialized by 'writeBuffer' from the provided
        /// input stream.
        BufferPtr readBuffer(std::istream &in, memory::MemoryPool *pool) {
            auto numBytes = read<int32_t>(in);
            auto buffer = AlignedBuffer::allocate<char>(numBytes, pool);
            auto rawBuffer = buffer->asMutable<char>();
            in.read(rawBuffer, numBytes);
            return buffer;
        }

        /// Deserializes a optional BufferPtr serialized by 'writeOptionalBuffer' from
        /// the provided input stream.
        BufferPtr readOptionalBuffer(std::istream &in, memory::MemoryPool *pool) {
            bool hasBuffer = read<bool>(in);
            if (hasBuffer) {
                return readBuffer(in, pool);
            }

            return nullptr;
        }

        template<TypeKind kind>
        VectorPtr createFlat(
            const TypePtr &type,
            vector_size_t size,
            pollux::memory::MemoryPool *pool,
            BufferPtr nulls,
            BufferPtr values,
            std::vector<BufferPtr> stringBuffers) {
            using T = typename TypeTraits<kind>::NativeType;

            return std::make_shared<FlatVector<T> >(
                pool,
                type,
                std::move(nulls),
                size,
                std::move(values),
                std::move(stringBuffers));
        }

        int32_t computeStringOffset(
            StringView value,
            const std::vector<BufferPtr> &stringBuffers) {
            int32_t offset = 0;
            for (const auto &buffer: stringBuffers) {
                auto start = buffer->as<char>();

                if (value.data() >= start && value.data() < start + buffer->size()) {
                    return (value.data() - start) + offset;
                }

                offset += buffer->size();
            }

            POLLUX_FAIL("String view points outside of the string buffers");
        }

        const char *computeStringPointer(
            int32_t offset,
            const std::vector<BufferPtr> &stringBuffers) {
            int32_t totalOffset = 0;
            for (const auto &buffer: stringBuffers) {
                auto size = buffer->size();
                if (offset >= totalOffset && offset < totalOffset + size) {
                    return buffer->as<char>() + offset - totalOffset;
                }

                totalOffset += buffer->size();
            }

            POLLUX_FAIL("String offset is outside of the string buffers: {}", offset);
        }

        void writeStringViews(
            vector_size_t size,
            const BufferPtr &strings,
            const std::vector<BufferPtr> &stringBuffers,
            std::ostream &out) {
            write<int32_t>(strings->size(), out);

            auto rawBytes = strings->as<char>();
            auto rawValues = strings->as<StringView>();
            for (auto i = 0; i < size; ++i) {
                auto stringView = rawValues[i];
                if (stringView.isInline()) {
                    out.write(rawBytes + i * sizeof(StringView), sizeof(StringView));
                } else {
                    // Size.
                    write<int32_t>(stringView.size(), out);

                    // 4 bytes of zeros for prefix. We'll fill this in appropriately when
                    // deserializing.
                    write<int32_t>(0, out);

                    // Offset.
                    auto offset = computeStringOffset(stringView, stringBuffers);
                    write<int64_t>(offset, out);
                }
            }
        }

        void restoreVectorStringViews(
            vector_size_t size,
            const BufferPtr &strings,
            const std::vector<BufferPtr> &stringBuffers) {
            auto rawBytes = strings->as<char>();
            auto rawValues = strings->asMutable<StringView>();
            for (auto i = 0; i < size; ++i) {
                auto value = rawValues[i];
                if (!value.isInline()) {
                    auto offset = *reinterpret_cast<const int64_t *>(
                        rawBytes + i * sizeof(StringView) + 8);
                    rawValues[i] =
                            StringView(computeStringPointer(offset, stringBuffers), value.size());
                }
            }
        }

        bool isVarcharOrVarbinary(const BaseVector &vector) {
            return vector.type_kind() == TypeKind::VARCHAR ||
                   vector.type_kind() == TypeKind::VARBINARY;
        }

        void writeFlatVector(const BaseVector &vector, std::ostream &out) {
            // Nulls buffer.
            writeOptionalBuffer(vector.nulls(), out);

            // Values buffer.
            if (isVarcharOrVarbinary(vector)) {
                const auto &values = vector.values();
                if (values) {
                    write<bool>(true, out);

                    const auto &stringBuffers =
                            vector.as_flat_vector<StringView>()->stringBuffers();
                    writeStringViews(vector.size(), values, stringBuffers, out);
                } else {
                    write<bool>(false, out);
                }
            } else {
                writeOptionalBuffer(vector.values(), out);
            }

            // String buffers.
            if (isVarcharOrVarbinary(vector)) {
                const auto &stringBuffers =
                        vector.as_flat_vector<StringView>()->stringBuffers();
                write<int32_t>(stringBuffers.size(), out);

                for (const auto &buffer: stringBuffers) {
                    writeBuffer(buffer, out);
                }
            }
        }

        VectorPtr readFlatVector(
            const TypePtr &type,
            vector_size_t size,
            std::istream &in,
            memory::MemoryPool *pool) {
            // Nulls buffer.
            BufferPtr nulls = readOptionalBuffer(in, pool);

            // Values buffer.
            BufferPtr values = readOptionalBuffer(in, pool);

            // String buffers.
            std::vector<BufferPtr> stringBuffers;
            if (type->is_varchar() || type->is_varbinary()) {
                int32_t numStringBuffers = read<int32_t>(in);
                for (auto i = 0; i < numStringBuffers; ++i) {
                    stringBuffers.push_back(readBuffer(in, pool));
                }

                // Update the pointers in the StringViews.
                if (values) {
                    restoreVectorStringViews(size, values, stringBuffers);
                }
            }

            return POLLUX_DYNAMIC_SCALAR_TYPE_DISPATCH_ALL(
                        createFlat, type->kind(), type, size, pool, nulls, values, stringBuffers);
        }

        template<TypeKind kind>
        void writeScalarConstant(const BaseVector &vector, std::ostream &out) {
            using T = typename TypeTraits<kind>::NativeType;

            auto value = vector.as<ConstantVector<T> >()->value_at(0);
            out.write(reinterpret_cast<const char *>(&value), sizeof(T));

            if constexpr (std::is_same_v<T, StringView>) {
                if (!value.isInline()) {
                    write<int32_t>(value.size(), out);
                    out.write(value.data(), value.size());
                }
            }
        }

        void writeConstantVector(const BaseVector &vector, std::ostream &out) {
            bool isNull = vector.is_null_at(0);
            write<bool>(isNull, out);

            if (isNull) {
                return;
            }

            auto baseVector = vector.value_vector();
            write<bool>(baseVector == nullptr, out);

            if (baseVector) {
                saveVector(*baseVector, out);
                write<int32_t>(vector.as<ConstantVector<ComplexType> >()->index(), out);
            } else {
                POLLUX_DYNAMIC_SCALAR_TYPE_DISPATCH_ALL(
                    writeScalarConstant, vector.type_kind(), vector, out);
            }
        }

        template<TypeKind kind>
        VectorPtr readConstant(
            const TypePtr &type,
            vector_size_t size,
            pollux::memory::MemoryPool *pool,
            std::istream &in) {
            using T = typename TypeTraits<kind>::NativeType;

            T value = read<T>(in);

            if constexpr (std::is_same_v<T, StringView>) {
                if (!value.isInline()) {
                    auto stringSize = read<int32_t>(in);
                    BufferPtr stringBuffer = AlignedBuffer::allocate<char>(stringSize, pool);
                    in.read(stringBuffer->template asMutable<char>(), stringSize);

                    return std::make_shared<ConstantVector<T> >(
                        pool,
                        size,
                        false,
                        type,
                        StringView(stringBuffer->as<char>(), stringSize));
                }
            }

            return std::make_shared<ConstantVector<T> >(
                pool, size, false, type, std::move(value));
        }

        VectorPtr readConstantVector(
            const TypePtr &type,
            vector_size_t size,
            std::istream &in,
            memory::MemoryPool *pool) {
            // Is-null flag.
            auto isNull = read<bool>(in);

            if (isNull) {
                return BaseVector::create_null_constant(type, size, pool);
            }

            bool scalar = read<bool>(in);
            if (scalar) {
                return POLLUX_DYNAMIC_SCALAR_TYPE_DISPATCH_ALL(
                            readConstant, type->kind(), type, size, pool, in);
            }

            auto baseVector = restoreVector(in, pool);
            auto baseIndex = read<int32_t>(in);

            return BaseVector::wrap_in_constant(size, baseIndex, baseVector);
        }

        void writeDictionaryVector(const BaseVector &vector, std::ostream &out) {
            // Nulls buffer.
            writeOptionalBuffer(vector.nulls(), out);

            // Indices buffer.
            writeBuffer(vector.wrap_info(), out);

            // Base vector.
            saveVector(*vector.value_vector(), out);
        }

        VectorPtr readDictionaryVector(
            const TypePtr & /*type*/,
            vector_size_t size,
            std::istream &in,
            memory::MemoryPool *pool) {
            // Nulls buffer.
            BufferPtr nulls = readOptionalBuffer(in, pool);

            // Indices buffer.
            BufferPtr indices = readBuffer(in, pool);

            // Base vector.
            auto baseVector = restoreVector(in, pool);

            return BaseVector::wrap_in_dictionary(nulls, indices, size, baseVector);
        }

        void writeRowVector(const BaseVector &vector, std::ostream &out) {
            // Nulls buffer.
            writeOptionalBuffer(vector.nulls(), out);

            const auto *row_vector = vector.as<RowVector>();
            POLLUX_CHECK_NOT_NULL(
                row_vector, "Expected a RowVector, got: {}", vector.toString());

            // Child vectors.
            auto numChildren = row_vector->childrenSize();
            write<int32_t>(numChildren, out);

            for (auto i = 0; i < numChildren; ++i) {
                const auto &child = row_vector->childAt(i);

                write<bool>(child != nullptr, out);
                if (child) {
                    saveVector(*child, out);
                }
            }
        }

        VectorPtr readRowVector(
            const TypePtr &type,
            vector_size_t size,
            std::istream &in,
            memory::MemoryPool *pool) {
            // Nulls buffer.
            BufferPtr nulls = readOptionalBuffer(in, pool);

            // Child vectors.
            auto numChildren = read<int32_t>(in);
            std::vector<VectorPtr> children;
            children.reserve(numChildren);
            for (auto i = 0; i < numChildren; ++i) {
                bool present = read<bool>(in);
                if (present) {
                    children.push_back(restoreVector(in, pool));
                } else {
                    children.push_back(nullptr);
                }
            }

            return std::make_shared<RowVector>(pool, type, nulls, size, children);
        }

        void writeArrayVector(const BaseVector &vector, std::ostream &out) {
            // Nulls buffer.
            writeOptionalBuffer(vector.nulls(), out);

            // Offsets and sizes.
            auto array_vector = vector.as<ArrayVector>();
            writeBuffer(array_vector->offsets(), out);
            writeBuffer(array_vector->sizes(), out);

            // Elements vector.
            saveVector(*array_vector->elements(), out);
        }

        VectorPtr readArrayVector(
            const TypePtr &type,
            vector_size_t size,
            std::istream &in,
            memory::MemoryPool *pool) {
            // Nulls buffer.
            BufferPtr nulls = readOptionalBuffer(in, pool);

            BufferPtr offsets = readBuffer(in, pool);
            BufferPtr sizes = readBuffer(in, pool);

            auto elements = restoreVector(in, pool);

            return std::make_shared<ArrayVector>(
                pool, type, nulls, size, offsets, sizes, elements);
        }

        void writeMapVector(const BaseVector &vector, std::ostream &out) {
            // Nulls buffer.
            writeOptionalBuffer(vector.nulls(), out);

            // Offsets and sizes.
            auto map_vector = vector.as<MapVector>();
            writeBuffer(map_vector->offsets(), out);
            writeBuffer(map_vector->sizes(), out);

            // Keys and values vectors.
            saveVector(*map_vector->mapKeys(), out);
            saveVector(*map_vector->mapValues(), out);
        }

        VectorPtr readMapVector(
            const TypePtr &type,
            vector_size_t size,
            std::istream &in,
            memory::MemoryPool *pool) {
            // Nulls buffer.
            BufferPtr nulls = readOptionalBuffer(in, pool);

            BufferPtr offsets = readBuffer(in, pool);
            BufferPtr sizes = readBuffer(in, pool);

            auto keys = restoreVector(in, pool);
            auto values = restoreVector(in, pool);

            return std::make_shared<MapVector>(
                pool, type, nulls, size, offsets, sizes, keys, values);
        }

        void writeLazyVector(const BaseVector &vector, std::ostream &out) {
            auto lazyVector = dynamic_cast<const LazyVector *>(&vector);
            // check if the vector was loaded.
            bool isLoaded = lazyVector->isLoaded();
            write<bool>(isLoaded, out);

            // loaded vector.
            if (isLoaded) {
                saveVector(*vector.loaded_vector(), out);
            }
        }

        // A thin layer over the loaded vector. Returns the loaded vector as-is but
        // throws an error attempting to load when there was no original loaded vector.
        // NOTE: since this loads the vector as-is, the user needs to be mindful that
        // this will not interact with the rowSet passed and hence should only be used
        // to reproduce failure and not to test out changes.
        class LoadedVectorShim : public VectorLoader {
        public:
            explicit LoadedVectorShim(VectorPtr vector) : vector_(vector) {
            }

            void loadInternal(
                RowSet /*rowSet*/,
                ValueHook * /*hook*/,
                vector_size_t resultSize,
                VectorPtr *result) override {
                POLLUX_CHECK(
                    vector_ != nullptr, "This lazy vector should not have been loaded.");
                *result = vector_;
                POLLUX_CHECK_EQ((*result)->size(), resultSize);
            }

        private:
            // Is nullptr if there was no loaded vector to load.
            VectorPtr vector_;
        };

        VectorPtr readLazyVector(
            const TypePtr &type,
            vector_size_t size,
            std::istream &in,
            memory::MemoryPool *pool) {
            // check if the vector was loaded.
            bool isLoaded = read<bool>(in);
            VectorPtr loaded_vector;
            // loaded vector.
            if (isLoaded) {
                loaded_vector = restoreVector(in, pool);
            }
            return std::make_shared<LazyVector>(
                pool, type, size, std::make_unique<LoadedVectorShim>(loaded_vector));
        }
    } // namespace

    void saveType(const TypePtr &type, std::ostream &out) {
        auto serialized = toJson(type->serialize());
        write(serialized, out);
    }

    TypePtr restoreType(std::istream &in) {
        static melon::once_flag kOnce;

        melon::call_once(kOnce, []() { Type::registerSerDe(); });

        auto serialized = read<std::string>(in);
        auto json = melon::parseJson(serialized);
        return ISerializable::deserialize<Type>(json);
    }

    void saveVector(const BaseVector &vector, std::ostream &out) {
        // Encoding.
        writeEncoding(vector.encoding(), out);

        // Type kind.
        saveType(vector.type(), out);

        // Vector size.
        write(vector.size(), out);

        switch (vector.encoding()) {
            case VectorEncoding::Simple::FLAT:
                writeFlatVector(vector, out);
                return;
            case VectorEncoding::Simple::CONSTANT:
                writeConstantVector(vector, out);
                return;
            case VectorEncoding::Simple::DICTIONARY:
                writeDictionaryVector(vector, out);
                return;
            case VectorEncoding::Simple::ROW:
                writeRowVector(vector, out);
                return;
            case VectorEncoding::Simple::ARRAY:
                writeArrayVector(vector, out);
                return;
            case VectorEncoding::Simple::MAP:
                writeMapVector(vector, out);
                return;
            case VectorEncoding::Simple::LAZY:
                writeLazyVector(vector, out);
                return;
            default:
                POLLUX_UNSUPPORTED(
                    "Unsupported encoding: {}", mapSimpleToName(vector.encoding()));
        }
    }

    void saveVectorToFile(const BaseVector *vector, const char *filePath) {
        std::ofstream outputFile(filePath, std::ofstream::binary);
        saveVector(*vector, outputFile);
        outputFile.close();
    }

    void saveStringToFile(const std::string &content, const char *filePath) {
        std::ofstream outputFile(filePath, std::ofstream::binary);
        outputFile.write(content.data(), content.size());
        outputFile.close();
    }

    VectorPtr restoreVector(std::istream &in, memory::MemoryPool *pool) {
        POLLUX_CHECK_NOT_NULL(pool);

        // Encoding.
        auto encoding = readEncoding(in);

        // Type kind.
        auto type = restoreType(in);

        // Vector size.
        auto size = read<int32_t>(in);

        switch (encoding) {
            case Encoding::kFlat:
                if (type->is_row()) {
                    return readRowVector(type, size, in, pool);
                } else if (type->is_array()) {
                    return readArrayVector(type, size, in, pool);
                } else if (type->is_map()) {
                    return readMapVector(type, size, in, pool);
                }
                return readFlatVector(type, size, in, pool);
            case Encoding::kConstant:
                return readConstantVector(type, size, in, pool);
            case Encoding::kDictionary:
                return readDictionaryVector(type, size, in, pool);
            case Encoding::kLazy:
                return readLazyVector(type, size, in, pool);
            default:
                POLLUX_UNREACHABLE();
        }
    }

    VectorPtr restoreVectorFromFile(
        const char *filePath,
        memory::MemoryPool *pool) {
        std::ifstream inputFile(filePath, std::ifstream::binary);
        POLLUX_CHECK(!inputFile.fail(), "Cannot open file: {}", filePath);

        auto result = restoreVector(inputFile, pool);
        inputFile.close();
        return result;
    }

    std::string restoreStringFromFile(const char *filePath) {
        std::ifstream inputFile(filePath, std::ifstream::binary);
        POLLUX_CHECK(!inputFile.fail(), "Cannot open file: {}", filePath);

        // Find out file size.
        auto begin = inputFile.tellg();
        inputFile.seekg(0, std::ios::end);
        auto end = inputFile.tellg();

        auto fileSize = end - begin;
        if (fileSize == 0) {
            return "";
        }

        // Read the file.
        std::string result;
        result.resize(fileSize);

        inputFile.seekg(begin);
        inputFile.read(result.data(), fileSize);
        inputFile.close();
        return result;
    }

    std::optional<std::string> generateFolderPath(
        const char *basePath,
        const char *prefix) {
        auto path = fmt::format("{}/pollux_{}_XXXXXX", basePath, prefix);
        auto createdPath = mkdtemp(path.data());
        if (createdPath == nullptr) {
            return std::nullopt;
        }
        return path;
    }

    void saveSelectivityVector(const SelectivityVector &rows, std::ostream &out) {
        auto range = rows.asRange();

        write<int32_t>(rows.size(), out);
        out.write(
            reinterpret_cast<const char *>(range.bits()), bits::nbytes(rows.size()));
    }

    SelectivityVector restoreSelectivityVector(std::istream &in) {
        auto size = read<int32_t>(in);

        const auto numBytes = bits::nbytes(size);

        // SelectivityVector::setFromBits reads whole words (8 bytes). Make sure it is
        // safe to access the "last" partially-populated word.
        std::vector<char> bits(bits::roundUp(numBytes, 8));

        in.read(bits.data(), numBytes);

        SelectivityVector rows(size);
        rows.setFromBits(reinterpret_cast<uint64_t *>(bits.data()), size);
        return rows;
    }

    void saveSelectivityVectorToFile(
        const SelectivityVector &rows,
        const char *filePath) {
        std::ofstream outputFile(filePath, std::ofstream::binary);
        saveSelectivityVector(rows, outputFile);
        outputFile.close();
    }

    SelectivityVector restoreSelectivityVectorFromFile(const char *filePath) {
        std::ifstream inputFile(filePath, std::ifstream::binary);
        POLLUX_CHECK(!inputFile.fail(), "Cannot open file: {}", filePath);
        auto result = restoreSelectivityVector(inputFile);
        inputFile.close();
        return result;
    }
} // namespace kumo::pollux

template<>
struct fmt::formatter<kumo::pollux::Encoding> : formatter<int> {
    auto format(kumo::pollux::Encoding s, format_context &ctx) const {
        return formatter<int>::format(static_cast<int>(s), ctx);
    }
};
